from torchvision import models
from PIL import Image
import matplotlib.pyplot as plt
import torch
import torchvision.transforms as T
import numpy as np

def decode_segmap(image,nc=21):
    label_colors=np.array([(0,0,0),
                           (128,0,0),(0,128,0),(128,128,0),(0,0,128),(128,0,128),
                           (0,128,128),(128,128,128),(64,0,0),(192,0,0),(64,128,0),
                           (192,128,0),(64,0,128),(192,0,128),(64,128,128),(192,128,128),
                           (0,64,0),(128,64,0),(0,192,0),(128,192,0),(0,64,128)])
    r=np.zeros_like(image).astype(np.uint8)
    g=np.zeros_like(image).astype(np.uint8)
    b=np.zeros_like(image).astype(np.uint8)
    for I in range(0,nc):
        idx=image==I
        r[idx]=label_colors[I,0]
        g[idx]=label_colors[I,1]
        b[idx]=label_colors[I,2]

    rgb=np.stack([r,g,b],axis=2)
    return rgb

def segment(net,path):
    img=Image.open(path).convert('RGB')
    plt.imshow(img)
    trf=T.Compose([T.Resize(256),
                   T.CenterCrop(224),
                   T.ToTensor(),
                   T.Normalize(mean=[0.485,0.456,0.406],
                               std=[0.229,0.224,0.225])])
    inp=trf(img).unsqueeze(0)
    out=net(inp)['out']
    om=torch.argmax(out.squeeze(),dim=0).detach().cpu().numpy()
    rgb=decode_segmap(om)
    plt.imshow(rgb)
    plt.axis('off')
    plt.show()

fcn=models.segmentation.fcn_resnet50(pretrained=True).eval()
img='tt.png'
segment(fcn,img)
